#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
图像浮雕处理器 - 处理图片上的文字，实现凹版和凸版效果
"""

import os
import argparse
import time
import shutil
from PIL import Image, ImageFilter, ImageOps
import numpy as np

class ImageReliefProcessor:
    """图像浮雕处理器类，用于处理图片上的文字，实现凹版和凸版效果"""
    
    def __init__(self):
        """初始化图像浮雕处理器"""
        # 定义主输出目录
        self.main_output_dir = "processed_images_organized"
    
    def get_output_subdirectory(self, input_path, organization_strategy='source', custom_name=None, effect_type=None):
        """
        根据输入路径和组织策略生成输出子目录
        
        参数:
            input_path: 输入文件或目录的路径
            organization_strategy: 组织策略，可选值：
                'source': 根据源文件名或目录名组织
                'timestamp': 根据时间戳组织
                'custom': 使用自定义名称组织
            custom_name: 当organization_strategy为'custom'时使用的自定义名称
            effect_type: 效果类型，'convex'（凸版）或'concave'（凹版）
            
        返回:
            str: 生成的子目录路径
        """
        # 根据效果类型添加标记
        effect_suffix = ''
        if effect_type:
            if effect_type.lower() == 'convex':
                effect_suffix = '_tu'  # 凸版添加'tu'标记
            elif effect_type.lower() == 'concave':
                effect_suffix = '_ao'  # 凹版添加'ao'标记
        
        if organization_strategy == 'timestamp':
            # 使用当前时间戳作为子目录名
            timestamp = time.strftime('%Y%m%d_%H%M%S')
            return os.path.join(self.main_output_dir, f"{timestamp}{effect_suffix}")
        elif organization_strategy == 'custom' and custom_name:
            # 使用自定义名称作为子目录名
            return os.path.join(self.main_output_dir, f"{custom_name}{effect_suffix}")
        else:
            # 默认策略：根据源文件名或目录名组织
            base_name = os.path.basename(input_path)
            # 移除扩展名（如果是文件）
            base_name = os.path.splitext(base_name)[0]
            return os.path.join(self.main_output_dir, f"{base_name}{effect_suffix}")
    
    def load_image(self, image_path):
        """
        加载图像文件
        
        参数:
            image_path: 图像文件路径
            
        返回:
            PIL.Image: 加载的图像对象
        """
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"图像文件不存在: {image_path}")
        
        return Image.open(image_path)
    
    def convert_to_grayscale(self, image):
        """
        将图像转换为灰度图
        
        参数:
            image: PIL.Image对象
            
        返回:
            PIL.Image: 灰度图对象
        """
        return image.convert('L')
    
    def create_convex_effect(self, image, depth=3):
        """
        创建凸版效果
        
        参数:
            image: PIL.Image对象（灰度图）
            depth: 浮雕深度，默认值为3
            
        返回:
            PIL.Image: 凸版效果的图像对象
        """
        # 创建浮雕滤镜
        kernel = [
            [-1, -1, -1],
            [-1, depth + 8, -1],
            [-1, -1, -1]
        ]
        
        # 应用浮雕滤镜
        convex_image = image.filter(ImageFilter.Kernel((3, 3), [item for sublist in kernel for item in sublist], 1, 128))
        return convex_image
    
    def create_concave_effect(self, image, depth=3):
        """
        创建凹版效果
        
        参数:
            image: PIL.Image对象（灰度图）
            depth: 浮雕深度，默认值为3
            
        返回:
            PIL.Image: 凹版效果的图像对象
        """
        # 先创建凸版效果，然后反转图像以获得凹版效果
        convex_image = self.create_convex_effect(image, depth)
        concave_image = ImageOps.invert(convex_image)
        return concave_image
    
    def enhance_edges(self, image, factor=1.5):
        """
        增强图像边缘，使浮雕效果更明显
        
        参数:
            image: PIL.Image对象
            factor: 边缘增强因子，默认值为1.5
            
        返回:
            PIL.Image: 边缘增强的图像对象
        """
        # 转换为numpy数组进行处理
        img_array = np.array(image)
        
        # 应用边缘增强
        edges = Image.fromarray(img_array).filter(ImageFilter.FIND_EDGES)
        edges_array = np.array(edges)
        
        # 合并原始图像和边缘图像
        enhanced_array = np.clip(img_array + factor * edges_array, 0, 255).astype(np.uint8)
        
        return Image.fromarray(enhanced_array)
    
    def process_image(self, image_path, output_path=None, effect_type='convex', depth=3, 
                     enhance_edges_flag=True, organization_strategy='source', custom_name=None, 
                     preserve_color=True):
        """
        处理图像并保存结果
        
        参数:
            image_path: 输入图像路径
            output_path: 输出图像路径（可选，如不提供则根据组织策略自动生成）
            effect_type: 效果类型，'convex'（凸版）或'concave'（凹版），默认值为'convex'
            depth: 浮雕深度，默认值为3
            enhance_edges_flag: 是否增强边缘，默认值为True
            organization_strategy: 组织策略，可选值：'source'（默认）、'timestamp'、'custom'
            custom_name: 当organization_strategy为'custom'时使用的自定义名称
            preserve_color: 是否保留原始颜色，默认值为True（保留原始背景色和文字阴影）
            
        返回:
            str: 保存的图像路径
        """
        # 加载图像
        image = self.load_image(image_path)
        
        if preserve_color and image.mode == 'RGB':
            # 保留原始颜色的处理方式
            # 转换为灰度图用于创建浮雕效果
            grayscale_image = self.convert_to_grayscale(image)
            
            # 创建浮雕效果
            if effect_type.lower() == 'convex':
                relief_image = self.create_convex_effect(grayscale_image, depth)
            elif effect_type.lower() == 'concave':
                relief_image = self.create_concave_effect(grayscale_image, depth)
            else:
                raise ValueError("效果类型必须是'convex'（凸版）或'concave'（凹版）")
            
            # 增强边缘
            if enhance_edges_flag:
                relief_image = self.enhance_edges(relief_image)
            
            # 将浮雕效果应用到原始彩色图像上
            # 将图像转换为numpy数组进行处理
            image_array = np.array(image)
            relief_array = np.array(relief_image)
            
            # 创建一个3通道的浮雕数组
            relief_3channel = np.stack((relief_array, relief_array, relief_array), axis=2)
            
            # 使用浮雕效果作为亮度调整，但保留原始色彩
            # 将RGB转换为HSV以更好地控制亮度
            from PIL import ImageEnhance
            
            # 创建一个亮度增强器，使用浮雕效果作为亮度映射
            enhancer = ImageEnhance.Brightness(image)
            processed_image = image.copy()
            
            # 这里我们使用原始图像作为基础，叠加浮雕效果
            # 这样可以保留原始的背景色和文字阴影
            processed_image = Image.blend(image, ImageOps.colorize(relief_image, 'black', 'white'), alpha=0.3)
            
            # 调整对比度以增强效果
            contrast_enhancer = ImageEnhance.Contrast(processed_image)
            processed_image = contrast_enhancer.enhance(1.2)
        else:
            # 原始的灰度处理方式
            # 转换为灰度图
            grayscale_image = self.convert_to_grayscale(image)
            
            # 创建浮雕效果
            if effect_type.lower() == 'convex':
                processed_image = self.create_convex_effect(grayscale_image, depth)
            elif effect_type.lower() == 'concave':
                processed_image = self.create_concave_effect(grayscale_image, depth)
            else:
                raise ValueError("效果类型必须是'convex'（凸版）或'concave'（凹版）")
            
            # 增强边缘
            if enhance_edges_flag:
                processed_image = self.enhance_edges(processed_image)
        
        # 如果未提供输出路径，则根据组织策略自动生成
        if output_path is None:
            # 获取子目录（包含效果类型标记）
            sub_dir = self.get_output_subdirectory(image_path, organization_strategy, custom_name, effect_type)
            # 获取原始文件名并添加效果类型后缀
            base_name = os.path.basename(image_path)
            name_without_ext, ext = os.path.splitext(base_name)
            output_filename = f"{name_without_ext}_{effect_type}{ext}"
            output_path = os.path.join(sub_dir, output_filename)
        
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # 保存处理后的图像
        processed_image.save(output_path)
        
        return output_path
    
    def process_directory(self, input_dir, output_dir=None, effect_type='convex', depth=3, 
                         enhance_edges_flag=True, organization_strategy='source', custom_name=None, 
                         preserve_color=True):
        """
        批量处理目录中的所有图像
        
        参数:
            input_dir: 输入目录路径
            output_dir: 输出目录路径（可选，如不提供则根据组织策略自动生成）
            effect_type: 效果类型，'convex'（凸版）或'concave'（凹版），默认值为'convex'
            depth: 浮雕深度，默认值为3
            enhance_edges_flag: 是否增强边缘，默认值为True
            organization_strategy: 组织策略，可选值：'source'（默认）、'timestamp'、'custom'
            custom_name: 当organization_strategy为'custom'时使用的自定义名称
            
        返回:
            list: 已处理的文件列表
        """
        if not os.path.exists(input_dir):
            raise FileNotFoundError(f"输入目录不存在: {input_dir}")
        
        # 如果未提供输出目录，则根据组织策略自动生成（包含效果类型标记）
        if output_dir is None:
            output_dir = self.get_output_subdirectory(input_dir, organization_strategy, custom_name, effect_type)
        
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        
        # 获取目录中的所有图像文件
        image_extensions = ['.png', '.jpg', '.jpeg', '.bmp', '.gif']
        processed_files = []
        
        for file_name in os.listdir(input_dir):
            file_ext = os.path.splitext(file_name)[1].lower()
            if file_ext in image_extensions:
                input_path = os.path.join(input_dir, file_name)
                output_path = os.path.join(output_dir, file_name)
                
                try:
                    self.process_image(input_path, output_path, effect_type, depth, enhance_edges_flag, 
                              organization_strategy, custom_name, preserve_color)
                    processed_files.append(file_name)
                    print(f"已处理: {file_name}")
                except Exception as e:
                    print(f"处理{file_name}时出错: {str(e)}")
        
        return processed_files


def main():
    """主函数，提供命令行接口"""
    parser = argparse.ArgumentParser(description='图像浮雕处理器 - 处理图片上的文字，实现凹版和凸版效果')
    
    # 输入选项（互斥组）
    input_group = parser.add_mutually_exclusive_group(required=True)
    input_group.add_argument('--image-path', type=str, help='输入图像文件路径')
    input_group.add_argument('--input-dir', type=str, help='输入图像目录路径（批量处理）')
    
    # 输出选项
    parser.add_argument('--output-path', type=str, help='输出图像文件路径（可选，如果不提供则自动生成）')
    parser.add_argument('--output-dir', type=str, help='输出图像目录路径（可选，如果不提供则自动生成）')
    
    # 组织选项
    parser.add_argument('--organization-strategy', type=str, choices=['source', 'timestamp', 'custom'], 
                        default='source', help='输出文件的组织策略：source（按源文件/目录名）、timestamp（按时间戳）、custom（自定义）')
    parser.add_argument('--custom-name', type=str, help='当organization-strategy为custom时使用的自定义名称')
    
    # 处理选项
    parser.add_argument('--effect-type', type=str, default='convex', 
                        help='效果类型：1（凸版）、2（凹版）、convex（凸版）或concave（凹版），默认值为convex')
    parser.add_argument('--depth', type=int, default=3, help='浮雕深度，默认值为3')
    parser.add_argument('--enhance-edges', action='store_true', help='增强边缘，使浮雕效果更明显')
    parser.add_argument('--preserve-color', action='store_true', default=True, help='保留原始颜色和阴影效果，默认启用')
    parser.add_argument('--no-preserve-color', action='store_false', dest='preserve_color', help='不保留原始颜色，使用灰度处理')
    
    args = parser.parse_args()
    
    # 验证参数
    if args.organization_strategy == 'custom' and not args.custom_name:
        parser.error('当指定--organization-strategy=custom时，必须同时指定--custom-name')
    
    # 转换数字效果类型为文本
    if args.effect_type == '1':
        effect_type = 'convex'  # 1表示凸版效果
    elif args.effect_type == '2':
        effect_type = 'concave'  # 2表示凹版效果
    else:
        effect_type = args.effect_type  # 保留原始值
    
    # 创建处理器实例
    processor = ImageReliefProcessor()
    
    # 处理图像
    if args.image_path:
        # 处理单个图像
        output_path = processor.process_image(
            args.image_path, 
            args.output_path, 
            effect_type, 
            args.depth, 
            args.enhance_edges,
            args.organization_strategy,
            args.custom_name,
            args.preserve_color
        )
        print(f"图像已成功处理并保存到: {output_path}")
    else:
        # 批量处理目录中的图像
        processed_files = processor.process_directory(
            args.input_dir, 
            args.output_dir, 
            effect_type, 
            args.depth, 
            args.enhance_edges,
            args.organization_strategy,
            args.custom_name,
            args.preserve_color
        )
        print(f"已成功处理{len(processed_files)}个图像文件")


if __name__ == "__main__":
    main()